# -*- coding: utf-8 -*-
"""
🚀 物联网设备唤醒延迟测试平台（高精度版）
功能：
  - 支持 60,000 个 TCP 长连接保活
  - 每 3 分钟错峰发送加密心跳
  - 每分钟随机唤醒一批设备（UDP）
  - 精确记录唤醒延迟（UDP 实际发送 → TCP 响应）
  - 使用 uvloop（Linux/macOS）加速
  - 线程安全访问共享状态
  - 结果导出 Excel（含平均、P95、P99）
"""

import asyncio
import socket
import struct
import time
import random
import string
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import datetime
from typing import Dict, Optional, Set
import pandas as pd
import threading

# ==================== 尝试使用 uvloop（仅 Linux/macOS）====================
try:
    import uvloop
    asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
    print("🚀 使用 uvloop 加速异步事件循环")
except ImportError:
    print("⚠️ 未安装 uvloop，可使用 'pip install uvloop'（仅支持 Linux/macOS）")
except RuntimeError as e:
    if "Windows" in str(e):
        print("⚠️ uvloop 不支持 Windows，使用原生 asyncio")
    else:
        print(f"⚠️ uvloop 初始化失败: {e}")

# ==================== 配置参数 ====================

# 设备 ID（DID）的固定前缀，用于标识设备类型或厂商
# 示例：生成的 DID 格式为 IOTGDD-000000-ABC123
DID_PREFIX = "IOTGDD"

# 模拟的总设备数量（DID 数量）
# 控制整个测试平台建立多少个 TCP 长连接
# 示例：100 表示模拟 100 个物联网设备
TOTAL_DID_COUNT = 60000

# 每批并发建立的 TCP 连接数量
# 避免一次性发起所有连接导致“连接风暴”或资源耗尽
# 示例：每批连接 100 个设备
BATCH_CONNECT_SIZE = 2000

# 批次之间的间隔时间（单位：秒）
# 让系统有时间稳定处理上一批连接，防止瞬时负载过高
# 示例：每 1 秒发起下一批连接
BATCH_CONNECT_INTERVAL = 0.5  # 秒

# TCP 服务器地址（用于设备保活、心跳、接收唤醒响应）
TCP_HOST = "120.26.136.5"
TCP_PORT = 50000

# TCP 心跳包发送间隔（单位：秒）
# 每个设备每隔此时间发送一次加密心跳，维持长连接不被断开
# 当前设置为 30 秒（注意：注释写的是“3分钟”，实际是 30 秒，建议统一）
# ⚠️ 注意：若想实现 3 分钟（180 秒），应改为 180
TCP_HEARTBEAT_INTERVAL = 180  # 实际为 30 秒

# UDP 服务器地址（用于发送设备唤醒指令）
# 唤醒包通过 UDP 发送，触发设备向 TCP 通道上报数据
UDP_HOST = "120.26.136.5"
UDP_PORT = 50001

# 每分钟通过 UDP 唤醒的设备数量
# 控制测试流量强度，避免瞬间唤醒过多设备造成服务端压力
# 示例：每分钟随机唤醒 100 个在线设备
UDP_BATCH_SIZE = 1000

# 唤醒任务执行间隔（单位：秒）
# 示例：30 表示每 30 秒唤醒一批设备；60 表示每分钟；300 表示每 5 分钟
WAKEUP_INTERVAL = 5  # 默认仍为 60 秒（每分钟一次）

# UDP 套接字操作超时时间（单位：秒）
# 发送或接收 UDP 包时，若超过此时间无响应则视为失败
UDP_TIMEOUT = 5

# 发送 UDP 唤醒包时使用的线程池大小
# 控制并发发送的线程数量，提高发送效率，但避免系统资源耗尽
# 示例：最多 50 个线程并行发送 UDP 包
UDP_THREADS = 100

# 加密密钥，用于生成设备通信的身份凭证（加密心跳内容）
# 必须与服务端保持一致，确保安全认证通过
ENCRYPTION_KEY = "RSIOTGWakeupKey@"

# 加密算法类型标识（传给 iPN_StringEnc 函数）
# 示例：256 可能表示 AES-256 或自定义高强度加密模式
ENCRYPTION_TYPE = 256

# 测试运行总时长（单位：分钟）
# 程序启动后持续运行指定分钟数，之后自动汇总结果并退出
# 示例：运行 3 分钟（适合调试）
# ⚠️ 注释中写的是“7天”，但值为 3，建议确认是否应为 7*24*60
RUN_MINUTES = 10*24*60  # 运行10天

# 结果输出文件名，包含时间戳以避免重复覆盖
# 数据将导出为 Excel 文件，包含延迟明细和统计摘要（平均、P95、P99）
# 示例：wakeup_latency_20250828_150000.xlsx
OUTPUT_FILE = f"wakeup_latency_{datetime.now().strftime('%Y%m%d_%H%M%S')}.xlsx"

# ==================== 全局状态与线程安全锁 ====================
# ✅ DID → UDP 实际发送时间（精确到发送线程）
pending_wakeup: Dict[str, float] = {}

# ✅ 延迟结果：DID → {DID, UDP_Send_Time, TCP_Recv_Time, Latency_ms}
latency_results: Dict[str, Dict] = {}

# ✅ 活跃客户端集合（线程安全）
active_clients: Set['TCPDeviceClient'] = set()
active_clients_lock = threading.Lock()

# ✅ 连接计数（线程安全）
connected_count = 0
total_connected = 0
connect_lock = threading.Lock()

# ✅ pending_wakeup 字典访问锁
pending_lock = threading.Lock()

# 导入加密函数
try:
    from string_enc import iPN_StringEnc
except ImportError:
    print("❌ 错误：请确保 string_enc.py 存在并包含 iPN_StringEnc 函数")
    exit(1)


# ==================== 工具函数 ====================
def generate_did_list() -> list:
    """生成 DID 列表：IOTGDD-000000-XXXXXX"""
    chars = string.ascii_uppercase
    end = ''.join(random.choices(chars, k=6))
    print(f"后缀: {end}")
    return [
        f"IOTGDD-{i:06d}-{end}"
        for i in range(TOTAL_DID_COUNT)
    ]


def create_udp_packet(did: str) -> bytes:
    """构建 UDP 唤醒包"""
    payload = did.encode('utf-8')
    msg_len = len(payload)
    header = struct.pack(">III", 0xFFFFFFA1, 1, msg_len)
    return header + payload


def send_udp_wakeup(did: str) -> bool:
    """
    发送单个 UDP 唤醒包
    ✅ 在发送成功后立即记录真实发送时间（精确延迟测量）
    """
    packet = create_udp_packet(did)
    try:
        with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as sock:
            sock.settimeout(UDP_TIMEOUT)
            send_time = time.time()  # ✅ 真实发送时间（在发送线程中）
            sock.sendto(packet, (UDP_HOST, UDP_PORT))

            # ✅ 只有发送成功才记录时间（线程安全）
            with pending_lock:
                pending_wakeup[did] = send_time

        return True
    except Exception:
        return False


# ==================== TCP 客户端 ====================
class TCPDeviceClient:
    def __init__(self, did: str):
        self.did = did
        _, self.encrypted_data = iPN_StringEnc(ENCRYPTION_KEY, did, ENCRYPTION_TYPE)
        if isinstance(self.encrypted_data, str):
            self.encrypted_data = self.encrypted_data.encode('utf-8')
        self.reader: Optional[asyncio.StreamReader] = None
        self.writer: Optional[asyncio.StreamWriter] = None
        self.connected = False

    async def connect(self):
        try:
            self.reader, self.writer = await asyncio.wait_for(
                asyncio.open_connection(TCP_HOST, TCP_PORT),
                timeout=10
            )
            self.connected = True

            # ✅ 加锁添加到活跃集合
            with active_clients_lock:
                active_clients.add(self)

            # 更新连接计数
            with connect_lock:
                global connected_count, total_connected
                connected_count += 1
                total_connected += 1
            if connected_count % 1000 == 0 or connected_count == TOTAL_DID_COUNT:
                now = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
                print(f"[{now}] ✅ 已连接 {connected_count}/{TOTAL_DID_COUNT}")

            # 启动心跳与监听
            asyncio.create_task(self._start_heartbeat())
            asyncio.create_task(self._listen())

        except Exception as e:
            # print(f"❌ 连接失败 {self.did}: {e}")
            pass

    async def _start_heartbeat(self):
        """错峰发送心跳，避免风暴"""
        initial_delay = random.uniform(0, TCP_HEARTBEAT_INTERVAL)
        await asyncio.sleep(initial_delay)

        while self.connected:
            if self.writer and not self.writer.is_closing():
                try:
                    self.writer.write(self.encrypted_data)
                    await self.writer.drain()
                except Exception:
                    self.connected = False
                    break
            await asyncio.sleep(TCP_HEARTBEAT_INTERVAL)

    async def _listen(self):
        """监听 TCP 响应"""
        try:
            while self.connected:
                data = await self.reader.read(1024)
                if not data:
                    break

                recv_time = time.time()
                recv_time_str = datetime.fromtimestamp(recv_time).strftime("%Y-%m-%d %H:%M:%S.%f")[:-3]

                with pending_lock:
                    if self.did in pending_wakeup:
                        send_time = pending_wakeup[self.did]
                        delay = recv_time - send_time
                        latency_results[self.did] = {
                            "DID": self.did,
                            "UDP_Send_Time": datetime.fromtimestamp(send_time).strftime("%Y-%m-%d %H:%M:%S.%f")[:-3],
                            "TCP_Recv_Time": recv_time_str,
                            "Latency_ms": round(delay * 1000, 2)
                        }
                        print(f"[{recv_time_str}] ✅ 响应 | {self.did} | 延迟: {delay * 1000:.2f} ms")
                        del pending_wakeup[self.did]
        except Exception:
            pass
        finally:
            self.connected = False
            # ✅ 加锁移除
            with active_clients_lock:
                active_clients.discard(self)
            if self.writer and not self.writer.is_closing():
                self.writer.close()
                try:
                    await self.writer.wait_closed()
                except:
                    pass


# ==================== UDP 唤醒任务（同步函数）====================
def trigger_udp_wakeup_batch():
    """
    每分钟触发一批 UDP 唤醒
    ✅ 不再预设时间，由 send_udp_wakeup 真实发送时记录时间
    """
    # ✅ 加锁读取活跃客户端
    with active_clients_lock:
        live_dids = [client.did for client in active_clients if client.connected]

    if not live_dids:
        print(f"[{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}] ⚠️  当前无活跃设备，跳过唤醒")
        return

    selected_dids = random.choices(live_dids, k=min(UDP_BATCH_SIZE, len(live_dids)))
    submit_time_str = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    print(f"[{submit_time_str}] ⏱️  提交 UDP 唤醒任务: 数量={len(selected_dids)}")

    # 多线程发送 UDP（真实发送时记录时间）
    with ThreadPoolExecutor(max_workers=UDP_THREADS) as executor:
        futures = [executor.submit(send_udp_wakeup, did) for did in selected_dids]
        success = sum(1 for f in as_completed(futures) if f.result())

    done_time_str = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    print(f"[{done_time_str}] 🚀 UDP 唤醒完成 | 提交: {len(selected_dids)} | 成功: {success}")


# ==================== 主程序 ====================
# ==================== 主程序 ====================
async def main():
    print(f"🚀 物联网设备唤醒延迟测试平台启动")
    print(f"  TCP 长连接数: {TOTAL_DID_COUNT}")
    print(f"  分批连接: {BATCH_CONNECT_SIZE} 个/批")
    print(f"  UDP 唤醒: 每 {WAKEUP_INTERVAL} 秒 {UDP_BATCH_SIZE} 个")
    print(f"  运行时长: {RUN_MINUTES} 分钟")
    print(f"  结果文件: {OUTPUT_FILE}")

    did_list = generate_did_list()
    tcp_clients = [TCPDeviceClient(did) for did in did_list]

    # 分批连接，避免连接风暴
    print("🟡 正在分批建立 TCP 连接...")
    for i in range(0, len(tcp_clients), BATCH_CONNECT_SIZE):
        batch = tcp_clients[i:i + BATCH_CONNECT_SIZE]
        for client in batch:
            asyncio.create_task(client.connect())
        await asyncio.sleep(BATCH_CONNECT_INTERVAL)

    print(f"✅ 已提交 {len(tcp_clients)} 个连接任务，正在后台连接...")

    # 主循环：按 WAKEUP_INTERVAL 定时唤醒
    start_time = time.time()
    end_time = start_time + RUN_MINUTES * 60
    next_wakeup_time = start_time + WAKEUP_INTERVAL  # 第一次唤醒时间

    try:
        while time.time() < end_time:
            # 计算还要等多久才到下次唤醒
            sleep_time = next_wakeup_time - time.time()
            if sleep_time > 0:
                await asyncio.sleep(sleep_time)

            # 执行唤醒任务
            trigger_udp_wakeup_batch()

            # 更新下次唤醒时间
            next_wakeup_time += WAKEUP_INTERVAL

    except KeyboardInterrupt:
        print("\n🛑 用户中断")
    finally:
        save_results()


def save_results():
    """保存结果到 Excel"""
    if not latency_results:
        print("⚠️  未记录到任何唤醒延迟数据")
        return

    df = pd.DataFrame(list(latency_results.values()))
    df = df.sort_values("Latency_ms")

    avg = df["Latency_ms"].mean()
    p95 = df["Latency_ms"].quantile(0.95)
    p99 = df["Latency_ms"].quantile(0.99)

    with pd.ExcelWriter(OUTPUT_FILE, engine='openpyxl') as writer:
        df.to_excel(writer, sheet_name='Latency', index=False)
        stats = pd.DataFrame({
            "Metric": ["Average (ms)", "P95 (ms)", "P99 (ms)", "Total Samples"],
            "Value": [round(avg, 2), round(p95, 2), round(p99, 2), len(df)]
        })
        stats.to_excel(writer, sheet_name='Summary', index=False)

    print(f"📊 结果已保存至: {OUTPUT_FILE}")
    print(f"📈 平均延迟: {avg:.2f}ms | P95: {p95:.2f}ms | P99: {p99:.2f}ms")


if __name__ == "__main__":
    asyncio.run(main())